import argparse
import pickle
import zarr
import numpy as np
import os

def pkl_to_zarr(pkl_path, zarr_path):
    with open(pkl_path, 'rb') as f:
        data = pickle.load(f)
    root = zarr.open(zarr_path, mode='w')
    for k, v in data.items():
        root.create_dataset(k, data=v)
    print(f"已将 {pkl_path} 转换为 {zarr_path}")

def zarr_to_pkl(zarr_path, pkl_path):
    root = zarr.open(zarr_path, mode='r')
    data = {}
    for k in root.keys():
        data[k] = np.array(root[k])
    with open(pkl_path, 'wb') as f:
        pickle.dump(data, f)
    print(f"已将 {zarr_path} 转换为 {pkl_path}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="pkl/zarr 互转工具")
    parser.add_argument("mode", choices=["p2z", "z2p"], help="转换模式：p2z（pkl转zarr），z2p（zarr转pkl）")
    parser.add_argument("from_path", type=str, help="源文件路径")
    parser.add_argument("to_path", type=str, help="目标文件路径")
    args = parser.parse_args()

    if args.mode == "p2z":
        pkl_to_zarr(args.from_path, args.to_path)
    elif args.mode == "z2p":
        zarr_to_pkl(args.from_path, args.to_path)